import numpy as np
from dataPrepare import X_train, y_train, X_test, y_test, y_test_one_hot
from NeuralNetwork import NeuralNetwork
import time

# 设置随机种子，使结果可复现
np.random.seed(42)

# 步骤三：初始化神经网络
# 输入层节点数 (784 = 28x28 像素)
input_nodes = 784
# 隐藏层节点数 (可调整的超参数)
hidden_nodes = 200
# 输出层节点数 (10个数字类别)
output_nodes = 10
# 学习率 (可调整的超参数)
learning_rate = 0.1

# 创建神经网络实例
nn = NeuralNetwork(input_nodes, hidden_nodes, output_nodes, learning_rate)

# 步骤四：训练神经网络
print("开始训练神经网络...")
start_time = time.time()

# 设置训练轮数
epochs = 5

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    
    # 遍历训练数据
    for i in range(len(X_train)):
        # 训练神经网络
        nn.train(X_train[i], y_train[i])
        
        # 每1000个样本打印一次进度
        if (i + 1) % 1000 == 0:
            print(f"  已处理 {i+1}/{len(X_train)} 个样本")
    
    # 每个epoch结束后，计算在测试集上的准确率
    accuracy = nn.evaluate(X_test, y_test)
    print(f"  Epoch {epoch+1} 结束，测试集准确率: {accuracy:.4f}")

training_time = time.time() - start_time
print(f"训练完成！总用时: {training_time:.2f} 秒")

# 步骤五：评估神经网络
print("\n开始评估神经网络...")

# 计算最终准确率
final_accuracy = nn.evaluate(X_test, y_test)
print(f"最终测试集准确率: {final_accuracy:.4f}")

# 步骤六：可视化一些预测结果（可选）
print("\n展示一些预测示例:")
import matplotlib.pyplot as plt

# 从测试集中随机选择几个样本
num_examples = 5
example_indices = np.random.randint(0, len(X_test), num_examples)

plt.figure(figsize=(15, 3))
for i, idx in enumerate(example_indices):
    # 获取样本
    sample = X_test[idx]
    true_label = y_test[idx]
    
    # 进行预测
    outputs = nn.query(sample)
    predicted_label = np.argmax(outputs)
    
    # 显示图像和预测结果
    plt.subplot(1, num_examples, i+1)
    plt.imshow(sample.reshape(28, 28), cmap='gray')
    plt.title(f"预测: {predicted_label}, 真实: {true_label}")
    plt.axis('off')

plt.tight_layout()
plt.savefig('prediction_examples.png')
print("预测示例已保存为 'prediction_examples.png'")

# 保存训练好的模型（可选）
import pickle

with open('mnist_model.pkl', 'wb') as f:
    pickle.dump(nn, f)
print("模型已保存为 'mnist_model.pkl'") 